/*
* Description: 磁盘信息获取
*/
#include "diskpart_report.h"
#include <iostream>
#include <iomanip>
#include "log.h"
#include "rpc_service/sendto_host.h"
#include <cstdint>
#include <map>
#include <winscard.h>
#include <sstream>

namespace {
    const std::string PHYSICAL_DRIVER_PATH = "\\\\.\\PhysicalDrive0";
    const int PARTITION_BUFFER = 4096;
    const unsigned long EFI_GUID = 3240784680; // EFI分区识别码
    const std::string EVENT_NAME = "SA_VM_EFI_PARTITION_INFO";
    std::map <int, std::string> PARTITION_STYLE_MAP {
        {0, "PARTITION_STYLE_MBR"},
        {1, "PARTITION_STYLE_GPT"},
        {2, "PARTITION_STYLE_RAW"}
    };
}

namespace DFX {

void DiskpartReport::Init(std::shared_ptr<WINCONTROL::SendToHost> sendtoHost)
{
    if (sendtoHost == nullptr) {
        LOG_ERR("[DFX - DiskpartReport] - sendtoHost is null!");
        return;
    }
    m_sendtoHost = sendtoHost;
    Start();
}

void DiskpartReport::Deinit()
{
    if (m_sendtoHost != nullptr) {
        m_sendtoHost.reset();
    }
}

void DiskpartReport::Start()
{
    DiskInfo diskInfo;
    GetDiskInfo(diskInfo);
    DiskInfoReport(diskInfo);
    Deinit();
}

void DiskpartReport::GetDiskInfo(DiskInfo &diskInfo)
{
    // 获取磁盘设备句柄
    HANDLE hDevice = CreateFile(PHYSICAL_DRIVER_PATH.c_str(), GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE,
        NULL, OPEN_EXISTING, 0, NULL);
    if (hDevice == INVALID_HANDLE_VALUE) {
        LOG_ERR("Failed to open disk!");
        return;
    }

    // 获取分区信息  IOCTL_DISK_GET_DRIVE_LAYOUT_EX:输入/输出控制码（IOCTL），用于获取磁盘驱动器的分区布局信息。
    char partitionBuffer[PARTITION_BUFFER]{};
    DWORD bytesReturned{};
    if (!DeviceIoControl(hDevice, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, NULL, 0,
        partitionBuffer, sizeof(partitionBuffer), &bytesReturned, NULL)) {
        LOG_ERR("Failed to get partition information!");
        CloseHandle(hDevice);
        return;
    }

    // 解析分区信息
    DRIVE_LAYOUT_INFORMATION_EX* layoutInfo = reinterpret_cast<DRIVE_LAYOUT_INFORMATION_EX*>(partitionBuffer);
    if (layoutInfo == nullptr) {
        LOG_ERR("layoutInfo is nullptr!");
        CloseHandle(hDevice);
        return;
    }
    PARTITION_INFORMATION_EX* partitionInfo = layoutInfo->PartitionEntry;
    int partitionCount = layoutInfo->PartitionCount;

    if (partitionInfo == nullptr) {
        LOG_ERR("partitionInfo is nullptr!");
        return;
    }
    for (int i = 0; i < partitionCount; i++) {
        if (partitionInfo[i].Gpt.PartitionType.Data1 == EFI_GUID) {
            diskInfo.partitionFormat = PARTITION_STYLE_MAP[partitionInfo[i].PartitionStyle];
            diskInfo.partitionId = partitionInfo[i].PartitionNumber;
            break;
        }
    }
 
    // 关闭设备句柄
    CloseHandle(hDevice);
}

void DiskpartReport::DiskInfoReport(DiskInfo &diskInfo)
{
    // 事件上报
    LOG_INFO("report diskinfomation");
    std::ostringstream reportStream;
    reportStream << "{\"partId\":\"" << diskInfo.partitionId << "\",\"format:\":\"" <<
        diskInfo.partitionFormat << "\"}";
    if (m_sendtoHost == nullptr) {
        LOG_ERR("[DFX - DiskpartReport] - sendtoHost is null!");
        return;
    }
    LOG_INFO("diskinfomation:{}", reportStream.str());
}
}
